import torch
import numpy as np
import config


# class ReplayBuffer:
#     def __init__(self, args):
#         self.s = np.zeros((10240,13790))
#         # self.s = np.zeros((args.batch_size, args.statefinal_dim))
#         self.a = np.zeros((10240, 1))
#         self.a_logprob = np.zeros((10240, 1))
#         self.r = np.zeros((10240, 1))
#         self.s_ = np.zeros((10240,13790))
#         # self.s_ = np.zeros((args.batch_size, args.statefinal_dim ))
#         self.dw = np.zeros((10240, 1))
#         self.done = np.zeros((10240, 1))
#         self.count = 0

#     def store(self, s, a, a_logprob, r, s_, dw, done):
#         self.s[self.count] = s
#         self.a[self.count] = a
#         self.a_logprob[self.count] = a_logprob
#         self.r[self.count] = r
#         self.s_[self.count] = s_
#         self.dw[self.count] = dw
#         self.done[self.count] = done
#         self.count += 1

#     def numpy_to_tensor(self):
#         s = torch.tensor(self.s, dtype=torch.float)
#         a = torch.tensor(self.a, dtype=torch.long)  # In discrete action space, 'a' needs to be torch.long
#         a_logprob = torch.tensor(self.a_logprob, dtype=torch.float)
#         r = torch.tensor(self.r, dtype=torch.float)
#         s_ = torch.tensor(self.s_, dtype=torch.float)
#         dw = torch.tensor(self.dw, dtype=torch.float)
#         done = torch.tensor(self.done, dtype=torch.float)

#         return s, a, a_logprob, r, s_, dw, done
import numpy as np
import torch
import argparse,copy
class ReplayBuffer:
    def __init__(self, args, action_dim=10):  # 假设 action_dim=10
        self.capacity = 10240  # 根据您的实现
        self.s = np.zeros((self.capacity, 13790))
        self.a = np.zeros((self.capacity, 1))
        self.a_logprob = np.zeros((self.capacity, 1))
        self.r = np.zeros((self.capacity, 1))
        self.s_ = np.zeros((self.capacity, 13790))
        self.dw = np.zeros((self.capacity, 1))
        self.done = np.zeros((self.capacity, 1))
        self.a_prob_initial = np.zeros((self.capacity, action_dim))  # 新增的数组
        self.count = 0

    def store(self, s, a, a_logprob, r, s_, dw, done, a_prob_initial):
        if self.count >= self.capacity:
            self.count = 0  # 循环覆盖
        self.s[self.count] = s
        self.a[self.count] = a
        self.a_logprob[self.count] = a_logprob
        self.r[self.count] = r
        self.s_[self.count] = s_
        self.dw[self.count] = dw
        self.done[self.count] = done
        self.a_prob_initial[self.count] = a_prob_initial  # 存储 a_prob_initial
        self.count += 1

    def numpy_to_tensor(self):
        # 只转换存储了数据的部分
        batch_size = self.count if self.count < self.capacity else self.capacity
        s = torch.tensor(self.s[:batch_size], dtype=torch.float).to(device)
        a = torch.tensor(self.a[:batch_size], dtype=torch.long).to(device)
        a_logprob = torch.tensor(self.a_logprob[:batch_size], dtype=torch.float).to(device)
        r = torch.tensor(self.r[:batch_size], dtype=torch.float).to(device)
        s_ = torch.tensor(self.s_[:batch_size], dtype=torch.float).to(device)
        dw = torch.tensor(self.dw[:batch_size], dtype=torch.float).to(device)
        done = torch.tensor(self.done[:batch_size], dtype=torch.float).to(device)
        a_prob_initial = torch.tensor(self.a_prob_initial[:batch_size], dtype=torch.float).to(device)  # 转换 a_prob_initial

        return s, a, a_logprob, r, s_, dw, done, a_prob_initial

    def __len__(self):
        return self.count
